﻿using AiBLSmartEdu.Module.AIMedicalAssistant.API.Services;
using AutoMapper;
using ChromaDB.Client;
using ChromaDB.Client.Models;
using Domain.Entities;
using FrameworkCore.Hubs;
using FrameworkCore.Interfaces;
using FrameworkCore.Repositories;
using Microsoft.AspNetCore.SignalR;
using Polly;
using System.Collections.Concurrent;
using System.Text;

#nullable disable

namespace AiBLSmartEdu.Module.AIMedicalAssistant.API.Hubs
{
    public class NotificationHub : Hub
    {
        private static readonly ConcurrentDictionary<string, string> Users = new();

        private readonly MedicalAssistantService _mediicalAssistantService;
        private readonly KnowledgeBaseService _knowledgeBaseService;
        private readonly IRepository<MedicalQuestionAnswerPairs> _medicalQuestionAnswersRepository;
        private readonly IMapper _mapper;
        private readonly EmbeddingGenerator _embeddingGenerator;
        private readonly OllamaService _ollamaService;
        private readonly ChromaClientService _chromaClientService;
        private readonly ICurrentUserService _currentUserService;

        public NotificationHub(
            MedicalAssistantService mediicalAssistantService, 
            KnowledgeBaseService knowledgeBaseService, 
            IRepository<MedicalQuestionAnswerPairs> medicalQuestionAnswersRepository, 
            IMapper mapper, 
            EmbeddingGenerator embeddingGenerator, 
            OllamaService ollamaService, 
            ChromaClientService chromaClientService, 
            ICurrentUserService currentUserService)
        {
            _mediicalAssistantService = mediicalAssistantService;
            _knowledgeBaseService = knowledgeBaseService;
            _medicalQuestionAnswersRepository = medicalQuestionAnswersRepository;
            _mapper = mapper;
            _embeddingGenerator = embeddingGenerator;
            _ollamaService = ollamaService;
            _chromaClientService = chromaClientService;
            _currentUserService = currentUserService;
        }

        /// <summary>
        /// 接收问题并发送答案
        /// </summary>
        /// <param name="question">要发送的问题</param>
        /// <param name="dialogueId">对话ID</param>
        /// <param name="knowledgeBaseId">知识库ID</param>
        public virtual async Task SendAIMedicalAssistantQuestionMessage(string question, string knowledgeBaseId, string dialogueId)
        {
            const int embedDimension = 384; // 与集合维度一致
            const int maxRetries = 3;
            var collectionClient = _chromaClientService.GetCollectionClient(knowledgeBaseId.ToString());

            // 1. 生成问题嵌入（实际应调用模型API）
            var queryEmbedding = await _embeddingGenerator.GenerateEmbeddingAsync(question, embedDimension);
            var queryMemory = new ReadOnlyMemory<float>(queryEmbedding);

            // 2. 执行向量查询
            var results = await Policy
                .Handle<HttpRequestException>()
                .RetryAsync(maxRetries)
                .ExecuteAsync(async () =>
                    await collectionClient.Query(
                        queryEmbeddings: new List<ReadOnlyMemory<float>> { queryMemory },
                        nResults: 100, // 增加返回结果数量
                        include: ChromaQueryInclude.Metadatas | ChromaQueryInclude.Distances
                    ));

            // 3. 构建优化后的上下文
            var context = BuildContext(results, maxLength: 100000);

            // 4. 生成增强型Prompt
            var prompt = $$"""
                请基于以下上下文用中文回答问题，遵循以下规则：
                1. 如果上下文无关，回答"该问题超出知识范围"
                2. 保持回答专业简洁
                3. 使用Markdown格式

                上下文：
                {{context}}

                问题：{{question}}
                答案：
                """;

            // 5. 获取模型回答
            var answer = await _ollamaService.GenerateResponseAsync(GetUserId(), prompt, "deepseek-r1");

            await _mediicalAssistantService.UpdateAnswerAsync(GetUserId(), question, answer, knowledgeBaseId, dialogueId);
            await Clients.Caller.SendAsync("ReceiveAIMedicalAssistantAnswerMessage", answer);

            string BuildContext(List<List<ChromaCollectionQueryEntry>> results, int maxLength)
            {
                var context = new StringBuilder();
                var entries = results.FirstOrDefault()
                    ?.OrderBy(e => e.Distance) // 按相似度排序
                    ?? Enumerable.Empty<ChromaCollectionQueryEntry>();

                foreach (var entry in entries)
                {
                    var content = entry.Metadata?
                        .GetValueOrDefault("content")?
                        .ToString()?
                        .Trim();

                    if (!string.IsNullOrEmpty(content))
                    {
                        // 检查是否还有空间添加新内容
                        if (context.Length + content.Length + 2 > maxLength) // +2 for "\n"
                        {
                            // 如果剩余空间不足，检查是否至少有一个完整段落
                            if (context.Length > 0) break;
                        }

                        context.AppendLine(content); // 直接添加内容，不加破折号
                    }
                }
                return context.ToString();
            }
        }

        /// <summary>
        /// 客户端连接时的逻辑。
        /// </summary>
        public override async Task OnConnectedAsync()
        {
            Users.TryAdd(Context.GetAccessToken(), Context.GetUserId());

            await base.OnConnectedAsync();
        }

        /// <summary>
        /// 获取用户链接Id
        /// </summary>
        /// <returns></returns>
        private string GetUserId()
        {
            Users.TryGetValue(Context.GetAccessToken(), out string? userId);
            return userId;
        }
    }
}
